(*
 *  MTASC - MotionTwin ActionScript2 Compiler
 *  Copyright (c)2004 Nicolas Cannasse
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

open ActionScript
open Expr
open SwfTypes
open ExtHashtbl
open ExtString
open ExtList
open GenSwfTools
open Extjava

let new_context ctx =
	{ ctx with
		idents = newHashTable [emptyJNup];
		locals = newHashTable [emptyJNup];
		stack_size = 0;
		ops = newList [];
		code_pos = 1;
		ident_count = 0;
		stack = 0;
		reg_count = 0;
		opt_push = false;
	}

(* The AS2 bytecode language requires jump offsets to be 16-bit. As a result,
   we have to eliminate overlength jumps. We do this by inserting "islands";
   i.e., the offending jump gets redirected to a new jump that we insert
   partway to the destination, and that jump goes the rest of the way (or to
   another jump, etc.). A second additional jump is inserted immediately before
   the island, which skips to the next instruction. Naturally, all other jump
   indices must then be adjusted. We endeavour to place jump islands at somewhat
   efficient locations (specifically, outside of nested function definitions),
   though there is no guarantee that we accomplish this with a minimal number of
   islands.

   Note that jump island insertion is not always possible. In such cases this
   algorithm will run forever. However, this basically never happens in
   practice. *)
(* This function requires tons of helper functions which are all very different
   from the rest of the code for class generation, so we stick them all inside
   this definition to keep it self-contained. *)
let eliminate_overlength_jumps =
	let maybe_get_jump_target act =
		match act with
			AJump target -> Some target
		| ACondJump target -> Some target
		| _ -> None
	in let get_jump_target act =
		match (maybe_get_jump_target act) with
			Some target -> target
		| None -> raise (Error "Not a jump instruction")
	in let set_jump_target acts i target =
		match (asActionType (asAction (listGet acts i))) with
                | (AJump _) -> listSet acts i (JAction (newAction (AJump target)))
		| (ACondJump _) -> listSet acts i (JAction (newAction (ACondJump target)))
		| _ -> raise (Error "Not a jump instruction")
	in let find_overlength_jump acts =
		(* Returns the index in "acts" of the first overlength jump, if any. *)
		(* We tail-recursively iterate over the actions and check each jump's size. *)
		let rec helper acts i =
			if i == (listLength acts) - 1 then
                            None (* end of array; no overlength jumps *)
			else
                            let curAct = asActionType (asAction (listGet acts i)) in
                            let maybe_target = maybe_get_jump_target curAct
                            in match maybe_target with
                                    None -> helper acts (i + 1) (* not a jump; continue *)
                            | Some target ->
                                    let size = jump_index_to_size acts i target in
                                    if size < -0x8000 || size > 0x7FFF
                                    (* Uncomment to stick jump islands all over the place for testing.
                                       (Must be larger than the largest instruction in the code.) *)
                                    (* if size < -0x106 || size > 0x100 *)
                                    then Some i (* overlength jump *)
                                    else helper acts (i + 1) (* not overlength; continue *)
		in helper acts 0
	in let fixup_target i target island_index island_length =
		(* Adjusts the instruction offset "target" for the instruction at index "i" to
		   account for an island having been inserted at "island_index" of length
		   "island_length". Note: given the choice, we jump to the beginning of
		   the island instead of the end, since we want to avoid putting islands inside
		   of nested functions. *)
		if i < island_index && (i + target + 1) > island_index then
			(* Jump goes forwards from before the island to after/into it; increase
			   by length of island. *)
			target + island_length
		else if i >= (island_index + island_length) && (i + target + 1) < (island_index + island_length) then
			(* Jump goes backwards from after the island to before/into it; decrease
			   by length of island (i.e., increase magnitude). *)
			target - island_length
		else
			(* Jump does not intersect with island; no change. *)
			target
	in let fixup_instruction_targets acts island_index island_length exempt_index =
		(* Iterates through an action sequence and fixes all instruction offsets
		   to account for new code having been inserted at island_index and of length
		   island_length, except for instructions inside it or at exempt_index. *)

                let helper i = 
                    if (i >= island_index && i < island_index + island_length) || i == exempt_index then function _ -> ()
                    else function
                            | AJump target ->
                                    listSet acts i (JAction (newAction (AJump (fixup_target i target island_index island_length))))
                            | ACondJump target ->
                                    listSet acts i (JAction (newAction (ACondJump (fixup_target i target island_index island_length))))
                            | AFunction f ->
                                    let oldCodeLen = asInt (getMField (JFuncDecl f) "codelen") in
                                    let newCodeLen = fixup_target i oldCodeLen island_index island_length in
                                    let setArgs = JArray ([JString "codelen" ; JInt newCodeLen]) in
                                    let functionDecl = asFuncDecl (setMField (JFuncDecl f) setArgs) in
                                    listSet acts i (JAction (newAction (AFunction functionDecl)))
                            | AFunction2 f ->
                                    let oldCodeLen = asInt (getMField (JFunc2Decl f) "codelen") in
                                    let newCodeLen = fixup_target i oldCodeLen island_index island_length in
                                    let setArgs = JArray ([JString "codelen" ; JInt newCodeLen]) in
                                    let function2Decl = asFunc2Decl (setMField (JFunc2Decl f) setArgs) in
                                    listSet acts i (JAction (newAction (AFunction2 function2Decl)))
                            | AWith size ->
                                    listSet acts i (JAction (newAction (AWith (fixup_target i size island_index island_length))))
                            | ATry t ->
                                    let tryLen = asInt (getMField (JTry t) "trylen") in
                                    let fixed_trylen = fixup_target i tryLen island_index island_length in
                                    let fixed_catchlen =
                                        match asIntOption (getMField (JTry t) "catchlen") with
                                            | None -> None
                                            | Some catchlen -> Some (fixup_target (i + fixed_trylen) catchlen island_index island_length) in
                                    let fixed_finallylen =
                                        match asIntOption (getMField (JTry t) "finallylen") with
                                            | None -> None
                                            | Some finallylen -> Some (fixup_target (i + fixed_trylen + (match fixed_catchlen with None -> 0 | Some i -> i)) finallylen island_index island_length) in
                                    listSet acts i (JAction (newAction (ATry (newTry (getTryStyleType t) fixed_trylen fixed_catchlen fixed_finallylen))))
                            | _ -> ()
                in
                for i = 0 to listLength acts do
                    helper i (asActionType (asAction (listGet acts i)))
                done
	in let get_function_target act =
		(* Returns the jump index needed to get to the end of the given function
		   instruction's definition, or 0 if not a function instruction. *)
		match act with
                | AFunction f -> asInt(getMField (JFuncDecl f) "codelen")
		| AFunction2 f -> asInt(getMField (JFunc2Decl f) "codelen")
		| _ -> 0
	in let jump_instruction_size =
		(* Constant value equal to the size of an unconditional jump instruction in
		   bytes. *)
		action_length (AJump 0)
	in let insert_jump_island acts i orig_target island_index =
		(* Inserts a jump island at "island_index" for the jump at "i" of length
		   "orig_target". *)
		(* "target" is relative to the instruction after "i", and "insert" inserts
		   _before_ the given index, so that's -1, but the hop jump has to come first,
		   so that's +1, so this is just island_index - i. *)
		let target = island_index - i in (* new target for the jump *)
		(* Compute the target value for the island's main jump. We will have gone
		   "target" instructions worth so far (so -target), but the number of
		   instructions to cover will have increased by 2 (so +2 if forward, -2 if
		   back). However, the target is relative to the instruction _after_ the new
		   jump (so -1). *)
		let island_target = orig_target - target - 1 + if orig_target > 0 then 2 else -2 in
		(* Now actually insert the jump island. *)

		listInsert acts island_index (JAction (newAction (AJump 1))); (* hop jump *)
		listInsert acts (island_index + 1) (JAction (newAction (AJump island_target))); (* main jump *)

		(* Now change the original jump. Note that if the jump goes backwards then
		   we've put the island before it so the index has changed. *)

		let new_i = if orig_target > 0 then i else i+2 in
		set_jump_target acts new_i target;

		(* Now iterate through the code and fix the targets of all other
		   instructions. *)

		fixup_instruction_targets acts island_index 2 new_i
	in let is_region_larger_than acts start stop compare_to =
		(* Checks whether or not the sum of action lengths in indices [start, end) is
		   greater than the given number. *)
		let rec helper acts curr stop acc compare_to =
			if acc > compare_to then true (* already larger *)
			else if curr >= stop then false (* done. never became larger so must not be *)
			else begin
                            let curAct = asActionType (asAction (listGet acts curr)) in
                            helper acts (curr+1) stop (acc+(action_length curAct)) compare_to
                        end
		in helper acts start stop 0 compare_to
	in let choose_nice_jump_island_index_in_range acts start stop is_forward_not_backward =
		(* Chooses an index (if any) in the range (start, end] where we can usefully
		   insert a jump island for a jump between "start" and "end" (in either
		   direction) without putting the island inside a nested function definition,
		   which would typically be mildly inefficient.

		   Specifically, this considers locations that are not inside a nested
		   function block and which would decrease the largest jump size involved.
		   From among those such points, it chooses the one closest to the halfway
		   point. This will not always result in a minimal number of jumps, but no
		   one will really care.

		   It is _not_ a requirement that "start" and "end" are in the same function,
		   and indeed there are cases where that property will not hold, b/c even
		   if we give up on a nice island we will still call this for future island
		   insertion, yet at that point some jumps in the code will span
		   function definitions. *)
		(* Note: curr here is the index to examine as a potential place _after_ which
		   to insert an island, which is different from the semantics above. *)
		let rec helper acts start stop curr best =
			if curr >= stop then best
			else
				(* Check if this index is the start of a function. If so, we skip past
				   it. *)
                                let curAct = asActionType (asAction (listGet acts curr)) in
				let skip = get_function_target curAct in
				if skip > 0 then
					(* This is a function. Skip to the end. *)
					helper acts start stop (curr + skip) best
				else
					(* Else we could safely insert an island here. See if that will actually
					   decrease the jump sizes. If we didn't bother with this, we'd get
					   into an infinite loop whenever there's a function that we can't
					   jump over, because we can always add useless jump islands at the
					   edges. *)
					let useful = if is_forward_not_backward then
							(* Jump goes from start to end. Redirecting it to the island's main
							   jump will leave out the instructions after the island but add in
							   the island's hop jump. Thus the jump size decreases iff the former
							   is larger than the latter. Further, the island's main jump gets to
							   leave out all instructions before the island and does not add in
							   any, so that decreases iff there is at least one instruction before
							   the island. *)
							(is_region_larger_than acts (curr+1) stop jump_instruction_size)
									&& (curr > start)
						else
							(* Jump goes from end to start. Redirecting it to the island's main
							   jump will leave out the instructions before the island but add in
							   the island's main jump. Thus the jump size decreases iff the
							   former is larger than the latter. Further, the island's main jump
							   gets to leave out all instructions after the island, but adds in
							   itself and the hop jump before it, so that decreases iff the
							   former is larger than the latter. *)
							(is_region_larger_than acts start (curr+1) jump_instruction_size) &&
								(is_region_larger_than acts (curr+1) (stop+1)
									(2*jump_instruction_size))
					in let best =
						if not useful then
							best
						else
							Some (match best with
								None -> curr+1
							| Some best ->
								if (abs ((stop + start)/2 - best)) >
									(abs ((stop + start)/2 - (curr+1)))
								then curr+1 else best)
					in helper acts start stop (curr+1) best
		in helper acts start stop start None
	in let choose_nice_jump_island_index acts i orig_target =
		(* Like choose_island_index_in_range, but specified in terms of the jump and its
		   target. *)
		if orig_target > 0 then
			choose_nice_jump_island_index_in_range acts i (i+orig_target+1) true
		else
			choose_nice_jump_island_index_in_range acts (i+orig_target+1) i false
	in let split_jump acts i =
		(* Splits the jump in "acts" at index "i" by inserting a jump island partway to
		   the destination. Returns true if it was able to place it outside of any
		   nested functions, or false if it had to place it within one or more. *)
                let curAct = asActionType(asAction (listGet acts i)) in
		let orig_target = get_jump_target curAct in (* original jump target *)
		let (nice, island_index) = match choose_nice_jump_island_index acts i orig_target with
			None ->
			(* There is no nice spot to put an island, so just put it in the middle.
			   (This happens if the jump goes across a function of size greater than
			   about 32 KiB.) *)
			(false, (i + (i + orig_target + 1))/2)
		| Some i ->
			(* Use that. *)
			(true, i)
		(* For simplicity, we only insert one jump island--regardless of how many are
		   necessary--and leave insertion of additional ones up to the later
		   iterations (since the one we insert will itself be an overlength jump of
		   smaller size). *)
		in insert_jump_island acts i orig_target island_index;
		nice
	in function acts ->
		(* This is now the actual definition of eliminate_overlength_jumps. *)
		let rec helper acts count bad_count =
			let index = find_overlength_jump acts
			in match index with
				None ->
				if count != 0 then prerr_endline ("Note: Jump island insertion successful with " ^ (string_of_int count) ^ " island(s) created" ^ (if bad_count = 0 then "." else " (at least " ^ (string_of_int bad_count) ^ " of them inefficient)."));
			| Some i ->
				if count = 0 then begin
					prerr_endline ("Note: Bytecode sequence contains one or more overlength jumps; attempting jump island insertion.");
					(* Flush stderr so that the user will always see the above message even if jump island insertion runs forever,
					   which can happen for pathological cases. *)
					flush stderr
				end;
				let nice = split_jump acts i in
				helper acts (count + 1) (if nice then bad_count else bad_count + 1)
		in helper acts 0 0		